#!/usr/bin/env python3
# Convert a coreos-assembler build into a "release.json"
# Originally from https://github.com/coreos/fedora-coreos-releng-automation/blob/main/coreos-meta-translator/trans.py
# See also https://github.com/coreos/fedora-coreos-tracker/blob/main/Design.md#release-streams


from argparse import ArgumentParser
import json
import os
import requests

from cosalib.builds import Builds

FCOS_STREAM_ENDPOINT = "https://builds.coreos.fedoraproject.org/prod/streams"


def ensure_dup(inp, out, inp_key, out_key):
    '''
    If the out dictionary does not contain a value for out_key update it
    to be equal to the inp dictionaries inp_key value, if it does exist
    ensure the values are equal between the two dictionaries
    '''
    inv = inp.get(inp_key)
    v = out.setdefault(out_key, inv)
    if v != inv:
        raise Exception(f"Input Files do not appear to be for the same release ({v} != {inv})")


def url_builder(stream, version, arch, path):
    return f"{args.stream_baseurl}/{stream}/builds/{version}/{arch}/{path}"


def get_extension(path, modifier, arch):
    return path.rsplit(f'{modifier}.{arch}')[1][1:]


parser = ArgumentParser()
parser.add_argument("--workdir", help="cosa workdir")
parser.add_argument("--build-id", help="build id")
parser.add_argument("--distro", help="Distro selects stream defaults such as baseurl and format", choices=['fcos', 'rhcos'])
parser.add_argument("--stream-name", help="Override the stream ID (default is derived from coreos-assembler)")
parser.add_argument("--stream-baseurl", help="Override prefix URL for stream content", default=FCOS_STREAM_ENDPOINT)
parser.add_argument("--output", help="Output to file; default is build directory")
parser.add_argument("--url", help="URL to a coreos-assembler meta.json", default=[], action='append')
parser.add_argument("--no-signatures", help="Disable signature references", action='store_true')
args = parser.parse_args()


def gather_buildmeta_from_workdir():
    builds = Builds()
    # default to latest build if not specified
    if args.build_id:
        buildid = args.build_id
    else:
        buildid = builds.get_latest()
    print(f"Creating release.json for build {buildid}")
    base_builddir = f"builds/{buildid}"
    arches = builds.get_build_arches(buildid)
    parsed_builds = []
    for arch in arches:
        with open(os.path.join(base_builddir, arch, "meta.json")) as f:
            parsed_builds.append(json.load(f))

    return (base_builddir, parsed_builds)


out = {}
parsed_builds = []

if len(args.url) == 0:
    # FIXME: Remove this once https://github.com/coreos/fedora-coreos-pipeline/ is ported
    # not to pass --workdir (it always uses `.` anyways)
    if args.workdir not in (None, '.'):
        os.chdir(args.workdir)

    (builddir, parsed_builds) = gather_buildmeta_from_workdir()
    # Default to writing into the builddir for now
    if args.output is None:
        args.output = os.path.join(builddir, "release.json")
else:
    for url in args.url:
        print(f"Downloading {url}...")
        r = requests.get(url)
        r.raise_for_status()
        parsed_builds.append(r.json())

# If any existing data, inherit it (if it's non-empty)
if os.path.exists(args.output) and os.stat(args.output).st_size > 0:
    with open(args.output, 'r') as w:
        out = json.load(w)
        print(f"Using existing release file {args.output}")


def get_floating_tag(rel, tags):
    found = ""
    for tag in tags:
        if rel not in tag:
            if found != "":
                raise f"multiple floating tags within: {tags}"
            found = tag
    if found == "":
        raise f"failed to find floating tag within: {tags}"
    return found


# Append the coreos-assembler build json `input_` to `out`, the target release stream.
def append_build(out, input_):
    arch = input_.get("coreos-assembler.basearch")

    ensure_dup(input_, out, "buildid", "release")
    streamnamesrc = None
    if args.stream_name:
        streamnamesrc = {'branch': args.stream_name}
    else:
        streamnamesrc = input_.get('coreos-assembler.container-config-git')
    ensure_dup(streamnamesrc, out, 'branch', 'stream')

    def artifact(i):
        base_url = url_builder(out.get('stream'), out.get('release'), arch, i.get('path'))
        sig = "{}.sig".format(base_url)
        if args.no_signatures:
            sig = ''
        return {
            "location": base_url,
            "signature": sig,
            "sha256": i.get("sha256"),
            "uncompressed-sha256": i.get("uncompressed-sha256")
        }

    print(f"{out['stream']} stream")
    print(f"  {arch} images:")
    # build the architectures dict
    arch_dict = {"media": {}}
    ensure_dup(input_, arch_dict, "ostree-commit", "commit")
    platforms = ["aliyun", "applehv", "aws", "azure", "azurestack", "digitalocean", "exoscale", "gcp", "hyperv", "ibmcloud", "kubevirt", "metal", "nutanix", "openstack", "powervs", "qemu", "virtualbox", "vmware", "vultr", "qemu-secex"]
    for platform in platforms:
        if input_.get("images", {}).get(platform, None) is not None:
            print(f"   - {platform}")
            i = input_.get("images").get(platform)
            ext = get_extension(i.get('path'), platform, arch)
            arch_dict['media'][platform] = {
                "artifacts": {
                    ext: {
                        "disk": artifact(i)
                    }
                }
            }
    # Aliyun/AWS specific additions
    for meta_key, cloud, image_field in ("aliyun", "aliyun", "id"), ("amis", "aws", "hvm"):
        if input_.get(meta_key, None) is not None:
            arch_dict["media"].setdefault(cloud, {}).setdefault("images", {})
            for cloud_dict in input_.get(meta_key):
                arch_dict["media"][cloud]["images"][cloud_dict["name"]] = {
                    "image": cloud_dict[image_field]
                }

    # IBMCloud/PowerVS specific additions
    for meta_key, cloud, object_field, bucket_field, url_field in \
        ("ibmcloud", "ibmcloud", "object", "bucket", "url"), \
            ("powervs", "powervs", "object", "bucket", "url"):
        if input_.get(meta_key, None) is not None:
            arch_dict["media"].setdefault(cloud, {}).setdefault("images", {})
            for cloud_dict in input_.get(meta_key):
                arch_dict["media"][cloud]["images"][cloud_dict["region"]] = {
                    "object": cloud_dict[object_field],
                    "bucket": cloud_dict[bucket_field],
                    "url": cloud_dict[url_field]
                }
    # IBM Secure Execution specific additions
    i = input_.get("images", {}).get("ignition-gpg-key", None)
    if i is not None:
        arch_dict["media"]["qemu-secex"]["ignition-gpg-key"] = artifact(i)

    # GCP specific additions
    if input_.get("gcp", None) is not None:
        arch_dict["media"].setdefault("gcp", {}).setdefault("image", {})
        arch_dict["media"]["gcp"]["image"].update(input_.get("gcp", {}))
        arch_dict["media"]["gcp"]["image"]["name"] = arch_dict["media"]["gcp"]["image"].pop("image")
        # remove the url as we haven't decided to expose that information publicly yet
        arch_dict["media"]["gcp"]["image"].pop("url")

    # KubeVirt specific additions: https://github.com/coreos/stream-metadata-go/pull/41
    if input_.get("kubevirt", None) is not None:
        arch_dict["media"].setdefault("kubevirt", {}).setdefault("image", {})
        # The `image` field uses a floating tag and the `digest-ref` field uses
        # a digest pullspec. See: https://github.com/coreos/stream-metadata-go/pull/46.
        tag = get_floating_tag(input_["buildid"], input_["kubevirt"]["tags"])
        arch_dict["media"]["kubevirt"]["image"] = {
            "image": input_["kubevirt"]["image"] + f":{tag}",
            "digest-ref": input_["kubevirt"]["image"] + "@" + input_["kubevirt"]["digest"],
        }

    # Azure: https://github.com/coreos/stream-metadata-go/issues/13
    inputaz = input_.get("azure")
    if inputaz is not None:
        rhcosext = arch_dict.setdefault("rhel-coreos-extensions", {})
        rhcosext["azure-disk"] = {
            "url": inputaz["url"]
        }

    # metal specific additions
    arch_dict["media"]["metal"] = arch_dict["media"].get("metal", {})
    arch_dict["media"]["metal"]["artifacts"] = arch_dict["media"]["metal"].get("artifacts", {})
    i = input_.get("images", {}).get("metal4k", None)
    if i is not None:
        # the 4k image is kinda weird; we want it at the same level as e.g.
        # the regular 512b image, which normally is under `raw.xz`
        ext = get_extension(i['path'], 'metal4k', arch)
        arch_dict["media"]["metal"]["artifacts"][f"4k.{ext}"] = {
            "disk": artifact(i)
        }
    i = input_.get("images", {}).get("iso", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"]["installer.iso"] = {
            "disk": artifact(i)
        }
    i = input_.get("images", {}).get("kernel", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"].setdefault("installer-pxe", {})["kernel"] = artifact(i)
    i = input_.get("images", {}).get("initramfs", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"].setdefault("installer-pxe", {})["initramfs"] = artifact(i)
    i = input_.get("images", {}).get("live-iso", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"]["iso"] = {
            "disk": artifact(i)
        }
    i = input_.get("images", {}).get("live-kernel", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"].setdefault("pxe", {})["kernel"] = artifact(i)
    i = input_.get("images", {}).get("live-initramfs", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"].setdefault("pxe", {})["initramfs"] = artifact(i)
    i = input_.get("images", {}).get("live-rootfs", None)
    if i is not None:
        arch_dict["media"]["metal"]["artifacts"].setdefault("pxe", {})["rootfs"] = artifact(i)

    # if architectures as a whole or the individual arch is empty just push our changes
    if out.get('architectures', None) is None or out['architectures'].get(arch, None) is None:
        oa = out.get('architectures', {})
        oa[arch] = arch_dict
        out['architectures'] = oa
    # else check media warning if key present, appending if not
    else:
        out_arch = out['architectures'][arch]
        for media_type, val in arch_dict.get('media').items():
            if media_type not in out_arch['media']:
                out['architectures'][arch]['media'].update({media_type: val})
            elif val == out_arch['media'][media_type]:
                continue
            else:
                raise Exception("differing content detected for media type '{}'".format(media_type))


for build in parsed_builds:
    append_build(out, build)

with open(args.output, 'w') as w:
    json.dump(out, w)
    print(f"Successfully wrote release file at {args.output}")
